import os
import sys
import argparse
import cv2
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from tqdm import tqdm

# 设置中文字体
plt.rcParams["font.sans-serif"] = ["SimHei", "Microsoft YaHei", "Arial Unicode MS"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def get_image_hash(image_path, hash_size=8):
    """
    计算图片的感知哈希值
    :param image_path: 图片路径
    :param hash_size: 哈希大小
    :return: 图片的二进制哈希值（numpy数组）
    """
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        return None
    
    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 缩放图片
    resized = cv2.resize(gray, (hash_size + 1, hash_size))
    
    # 计算差异
    diff = resized[:, 1:] > resized[:, :-1]
    
    # 生成哈希值并返回
    return np.packbits(diff)

def hamming_distance(hash1, hash2):
    """
    计算两个哈希值之间的汉明距离
    :param hash1: 第一个哈希值（numpy数组）
    :param hash2: 第二个哈希值（numpy数组）
    :return: 汉明距离
    """
    return np.count_nonzero(hash1 != hash2)

def deduplicate_images(input_dir, dry_run=False, hash_threshold=3, visualize=False):
    """
    对目录中的图片进行去重
    :param input_dir: 输入图片目录
    :param dry_run: 仅显示重复图片，不删除
    :param hash_threshold: 哈希阈值，汉明距离小于等于此值的图片被认为是重复的
    """
    # 确保输入目录存在
    if not os.path.exists(input_dir):
        print(f"错误：输入目录不存在: {input_dir}")
        return
    
    # 记录哈希值和对应的文件路径
    hash_list = []
    file_list = []
    
    # 用于可视化的数据
    duplicate_count = 0
    non_duplicate_count = 0
    hamming_distances = []
    
    # 用于按目录统计非重复图片数量
    dir_counts = {}

    
    # 遍历目录中的所有图片文件
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif']
    
    print(f"开始扫描目录: {input_dir}")
    print(f"哈希阈值设置为: {hash_threshold}")
    
    # 先统计总文件数
    total_files = 0
    for root, dirs, files in os.walk(input_dir):
        for file in files:
            if any(file.lower().endswith(ext) for ext in image_extensions):
                total_files += 1
    
    # 使用tqdm显示进度
    processed_files = 0
    duplicate_count = 0
    with tqdm(total=total_files, desc=f'已处理: 0, 重复: 0', ncols=80) as pbar:
        for root, dirs, files in os.walk(input_dir):
            for file in files:
                if any(file.lower().endswith(ext) for ext in image_extensions):
                    file_path = os.path.join(root, file)
                    processed_files += 1
                    pbar.set_description(f'已处理: {processed_files}, 重复: {duplicate_count}')
                    pbar.update(1)
                
                # 计算图片哈希值
                img_hash = get_image_hash(file_path)
                
                if img_hash is None:
                    print(f"警告：无法读取图片: {file_path}")
                    continue
                
                # 检查是否重复
                is_duplicate = False
                duplicate_with = None
                
                for i, existing_hash in enumerate(hash_list):
                    distance = hamming_distance(img_hash, existing_hash)
                    if distance <= hash_threshold:
                        is_duplicate = True
                        duplicate_with = file_list[i]
                        break
                
                if is_duplicate:
                    duplicate_count += 1
                    pbar.set_description(f'已处理: {processed_files}, 重复: {duplicate_count}')
                    hamming_distances.append(distance)
                    if not dry_run:
                        try:
                            os.remove(file_path)
                        except Exception as e:
                            print(f"  无法删除图片: {file_path}, 错误: {str(e)}")
                else:
                    hash_list.append(img_hash)
                    file_list.append(file_path)
                    non_duplicate_count += 1
                    
                    # 更新目录计数
                    dir_path = os.path.dirname(file_path)
                    if dir_path not in dir_counts:
                        dir_counts[dir_path] = 0
                    dir_counts[dir_path] += 1
    
    print(f"扫描完成。共发现 {duplicate_count} 张重复图片，{non_duplicate_count} 张非重复图片。")
    
    # 打印每个目录的非重复图片数量
    print("\n各目录非重复图片数量：")
    for dir_path, count in dir_counts.items():
        # 计算相对路径以便显示
        rel_dir = os.path.relpath(dir_path, input_dir)
        print(f"  {rel_dir}: {count} 张")
    
    # 生成可视化图表
    if visualize:
        visualize_results(duplicate_count, non_duplicate_count, hamming_distances, input_dir)

def visualize_results(duplicate_count, non_duplicate_count, hamming_distances, input_dir):
    """
    生成重复图片和非重复图片的对比图
    :param duplicate_count: 重复图片数量
    :param non_duplicate_count: 非重复图片数量
    :param hamming_distances: 重复图片的汉明距离列表
    :param input_dir: 输入目录
    """
    # 创建一个2x1的图表布局
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
    
    # 饼图：显示重复和非重复图片的比例
    labels = ['重复图片', '非重复图片']
    sizes = [duplicate_count, non_duplicate_count]
    colors = ['#ff9999', '#66b3ff']
    explode = (0.1, 0)
    
    ax1.pie(sizes, explode=explode, labels=labels, colors=colors, autopct='%1.1f%%',
            shadow=True, startangle=90)
    ax1.axis('equal')  # 确保饼图是圆的
    ax1.set_title('重复图片与非重复图片比例')
    
    # 直方图：显示重复图片的汉明距离分布
    if hamming_distances:
        ax2.hist(hamming_distances, bins=max(10, len(set(hamming_distances))), color='#99ff99', alpha=0.7)
        ax2.axvline(np.mean(hamming_distances), color='r', linestyle='dashed', linewidth=2, label=f'平均距离: {np.mean(hamming_distances):.2f}')
        ax2.set_xlabel('汉明距离')
        ax2.set_ylabel('图片数量')
        ax2.set_title('重复图片的汉明距离分布')
        ax2.legend()
    else:
        ax2.text(0.5, 0.5, '无重复图片', ha='center', va='center', transform=ax2.transAxes)
        ax2.axis('off')
    
    # 添加总标题
    plt.suptitle(f'图片去重结果分析', fontsize=16)
    
    # 调整布局
    plt.tight_layout(rect=[0, 0, 1, 0.95])
    
    # 保存图表
    output_path = os.path.join(os.path.dirname(input_dir), 'deduplication_results.png')
    plt.savefig(output_path, dpi=300)
    print(f"可视化图表已保存至: {output_path}")
    
    # 显示图表
    plt.show()

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='图片去重工具')
    parser.add_argument('-i', '--input', default='./output', help='输入图片目录（默认：./output）')
    parser.add_argument('-n', '--dry-run', action='store_true', help='仅显示重复图片，不删除')
    parser.add_argument('-t', '--threshold', type=int, default=5, help='哈希阈值（汉明距离），默认值：5')
    parser.add_argument('-v', '--visualize', action='store_true', help='生成重复图片和非重复图片的对比图')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 执行去重
    deduplicate_images(args.input, args.dry_run, args.threshold, args.visualize)

if __name__ == "__main__":
    main()